import concurrent.futures
import logging
from concurrent.futures import as_completed
from typing import List, Union
import copy

import dspy
dspy.dsp.settings.show_guidelines = False

from .callback import BaseCallbackHandler
from .storm_dataclass import StormInformationTable, StormArticle
from ...interface import ArticleGenerationModule, Information
from ...utils import ArticleTextProcessing

from .savetemp import savetemp


class StormArticleGenerationModule(ArticleGenerationModule):
    """
    The interface for article generation stage. Given topic, collected information from
    knowledge curation stage, generated outline from outline generation stage,
    """

    def __init__(
        self,
        article_gen_lm=Union[dspy.dsp.LM, dspy.dsp.HFModel],
        retrieve_top_k: int = 5,
        max_thread_num: int = 10,
    ):
        super().__init__()
        self.retrieve_top_k = retrieve_top_k
        self.article_gen_lm = article_gen_lm
        self.max_thread_num = max_thread_num
        self.section_gen = ConvToSection(engine=self.article_gen_lm)

    def generate_section(
        self, topic, section_name, information_table, section_outline, section_query, main_title
    ):
        import threading
        import datetime
        thread_id = threading.get_ident()
        savetemp("article_generation/thread_id.txt", section_name+' '+str(thread_id)+' '+str(datetime.datetime.now()))

        collected_info: List[Information] = []
        if information_table is not None:
            collected_info = information_table.retrieve_information(
                queries=section_query, search_top_k=self.retrieve_top_k
            )

        savetemp("article_generation/collected_info.txt", collected_info)

        output = self.section_gen(
            topic=topic,
            outline=section_outline,
            section=section_name,
            collected_info=collected_info,
        )
        # output_section = f"#{section_name}\n123\n"

        savetemp("article_generation/generate_section_output.txt", section_name+'\n'+output.section)
        # savetemp("article_generation/generate_section_output.txt", section_name+'\n'+output_section)

        section_level = len(section_query[0].split("::"))

        #change
        section_query = main_title + "::" + section_query[0]
        parent_section_name = section_query.split("::")[-2]
        
        import re
        def rewrite_section(section):
            lines = section.splitlines()
            if len(lines) <= 0:
                return section
            # 去除序号
            # lines[0] = re.sub(r"^(#+)\s*\d+(\.\d+)*\s*", r"\1 ", lines[0])
            # section_level个 #
            new_title = "#" * section_level + section_name
            lines.insert(0,new_title)
            # 替换#为** 去除references
            for i in range(2,len(lines)):
                if lines[i].strip("#").strip("-").strip() == "":
                    lines[i] = ""
                if lines[i].startswith("#"):
                    lines[i] = "**" + lines[i].strip("#").strip() + "**"
                # if start with [1-9] 
                if re.match(r'^\[\d+\]', lines[i].strip("*").strip()):
                    lines[i] = ""
                if re.match(r'(?i)^references?', lines[i].strip("*").strip()):
                    lines[i] = ""
                    
            return '\n'.join(line for line in lines if line.strip())
        
        section_content = rewrite_section(output.section)
        # section_content = rewrite_section(output_section)

        savetemp("article_generation/generate_section_output_rewrite.txt", section_name+'\n'+section_content)

        return {
            "section_name": section_name,
            "section_content": section_content,
            "collected_info": collected_info,
            #change
            "parent_section_name": parent_section_name,
        }

    def generate_article(
        self,
        topic: str,
        information_table: StormInformationTable,
        article_with_outline: StormArticle,
        callback_handler: BaseCallbackHandler = None,
    ) -> StormArticle:
        """
        Generate article for the topic based on the information table and article outline.

        Args:
            topic (str): The topic of the article.
            information_table (StormInformationTable): The information table containing the collected information.
            article_with_outline (StormArticle): The article with specified outline.
            callback_handler (BaseCallbackHandler): An optional callback handler that can be used to trigger
                custom callbacks at various stages of the article generation process. Defaults to None.
        """
        information_table.prepare_table_for_retrieval()

        savetemp("article_generation/raw_format_article_with_outline", article_with_outline)

        if article_with_outline is None:
            article_with_outline = StormArticle(topic_name=topic)

        sections_to_write = article_with_outline.get_first_level_section_names()

        section_output_dict_collection = []
        if len(sections_to_write) == 0:
            logging.error(
                f"No outline for {topic}. Will directly search with the topic."
            )
            section_output_dict = self.generate_section(
                topic=topic,
                section_name=topic,
                information_table=information_table,
                section_outline="",
                section_query=[topic],
                main_title=topic,
            )
            section_output_dict_collection = [section_output_dict]
        else:
            print("max thread num:", self.max_thread_num)
            with concurrent.futures.ThreadPoolExecutor(
                max_workers=self.max_thread_num
            ) as executor:
                future_to_sec_title = {}

                if len(sections_to_write) == 1:
                    main_title = sections_to_write[0]
                else:
                    main_title = topic

                # 所有一级标题
                for section_title in sections_to_write:

                    # # We don't want to write a separate introduction section.
                    # if section_title.lower().strip() == "introduction":
                    #     continue
                    #     # We don't want to write a separate conclusion section.
                    # if section_title.lower().strip().startswith(
                    #     "conclusion"
                    # ) or section_title.lower().strip().startswith("summary"):
                    #     continue
                    section_query = article_with_outline.get_outline_as_list(
                        root_section_name=section_title, add_hashtags=False
                    )
                    if len(section_query) == 0 or section_query[0].strip() == "":
                        continue
                    queries_with_hashtags = article_with_outline.get_outline_as_list(
                        root_section_name=section_title, add_hashtags=True
                    )
                    section_outline = "\n".join(queries_with_hashtags)

                    savetemp("article_generation/section_title.txt", section_title)
                    savetemp("article_generation/section_query.txt", section_query)
                    savetemp("article_generation/section_outline.txt", section_outline)


                    # change
                    from .gen_sub_title import get_sub_query
                    # 当前一级标题下的所有子标题 不限级数
                    sub_titles, sub_querys = get_sub_query(queries_with_hashtags)

                    print(f"sub title num:{len(sub_titles)}")

                    from tqdm import tqdm
                    
                    for sub_title,sub_query in tqdm(zip(sub_titles, sub_querys), total=len(sub_titles), desc="Generating sections"):

                        savetemp("article_generation/sub_title.txt", sub_title)
                        savetemp("article_generation/sub_query.txt", sub_query)

                        future_to_sec_title[
                            executor.submit(
                                self.generate_section,
                                topic = topic,
                                section_name = sub_title,
                                information_table = information_table,
                                section_outline = None,
                                section_query = [sub_query],
                                main_title = main_title,
                            )
                        ] = sub_title

                for future in as_completed(future_to_sec_title):
                    section_output_dict_collection.append(future.result())
        # change
        article = copy.deepcopy(article_with_outline)
        # article = StormArticle(topic_name=topic)

        # savetemp
        savetemp("article_generation/article.txt", article.root.content)

        for section_output_dict in section_output_dict_collection:
            
            savetemp("article_generation/save_process_section_name.txt","update section:"+section_output_dict["section_name"])
            savetemp("article_generation/save_process_parent_section_name.txt","update section:"+section_output_dict["parent_section_name"])

            article.update_section(
                # change
                # parent_section_name=topic,
                parent_section_name=section_output_dict["parent_section_name"],
                current_section_content=section_output_dict["section_content"],
                current_section_info_list=section_output_dict["collected_info"],
            )
        article.post_processing()
        return article


class ConvToSection(dspy.Module):
    """Use the information collected from the information-seeking conversation to write a section."""

    def __init__(self, engine: Union[dspy.dsp.LM, dspy.dsp.HFModel]):
        super().__init__()
        #change
        # self.write_section = dspy.Predict(WriteSection)
        self.write_section = dspy.Predict(WriteSection_sub_by_sub)
        self.engine = engine

    def forward(
        self, topic: str, outline: str, section: str, collected_info: List[Information]
    ):
        info = ""
        for idx, storm_info in enumerate(collected_info):
            info += f"[{idx + 1}]\n" + "\n".join(storm_info.snippets)
            info += "\n\n"

        # change 50000->12000->100
        info = ArticleTextProcessing.limit_word_count_preserve_newline(info, 20000)

        # savetemp
        savetemp("article_generation/info.txt", info)

        with dspy.settings.context(lm=self.engine):
            #change
            section = ArticleTextProcessing.clean_up_section(
                # self.write_section(topic=topic, section=section, word_limit="400", outline=outline, info=info).output
                self.write_section(topic=topic, section=section, info=info).output
            )

        return dspy.Prediction(section=section)


class WriteSection(dspy.Signature):
    """
    Write a section for a structured long literature review based on the collected information.

    Here is the format of your writing:
        1. Use "#" Title" to indicate section title, "##" Title" to indicate subsection title, "###" Title" to indicate subsubsection title, and so on.
        2. Use [1], [2], ..., [n] in line (for example, "The capital of the United States is Washington, D.C.[1][3].").
        You DO NOT need to include a References or Sources section to list the sources at the end.
    """

    info = dspy.InputField(prefix="The collected information:\n", format=str)
    topic = dspy.InputField(prefix="The topic of the page: ", format=str)

    #change
    outline = dspy.InputField(prefix="The outline you need to follow: ",format=str)

    section = dspy.InputField(prefix="The section you need to write: ", format=str)
    output = dspy.OutputField(
        prefix="Write the section with proper inline citations (Start your writing with # section title. Don't include the page title or try to write other sections):\n",
        format=str,
    )

# # change prompt
# class WriteSection_sub_by_sub(dspy.Signature):
#     """
#     Instruction: You are writing an overall and structured literature review based on the collected references.
#     Quirements:
#         1.Use the given section title as your title, ensuring the number of '#' symbols remains unchanged to indicate the title level.
#         DO NOT modify the section title.
#         DO NOT change the number of '#'.
#         2. When writing sentences based on specific papers, cite the paper using a numbered reference format [X], where X is the number corresponding to the paper in the referencelist at the end of the document.
#     Citation Guidelines:
#         1.Summarizing Research: Cite sources when summarizing existing literature.
#         2.Using Specific Concepts or Data: Provide citations when discussing specific theories, models, or data.
#         3.Using Established Methods: Cite the creators of methodologies you employ in yourliterature review.
#         4.Supporting Arguments: Cite sources that back up your conclusions and arguments
#     """

#     topic = dspy.InputField(prefix="The topic of the literature review: ", format=str)
#     section = dspy.InputField(prefix="Now, you need to write the content for a section: ", format=str)
#     word_limit = dspy.InputField(prefix="The section should contain words no more than: ", format=str)
#     # outline = dspy.InputField(prefix="The outline of this literature review is: ", format=str)
#     info = dspy.InputField(prefix="The following is a list of references:\n", format=str)
#     output = dspy.OutputField(
#     prefix="You DO NOT need to include a References or Sources section to list the sources at the end.Only return the content with words that you write for the section without any other information:\n",
#     format=str,
#     )

class WriteSection_sub_by_sub(dspy.Signature):
    """
### Task instruction
Your task is to generate a section of an overall and structured literature review answer based on the provided reference and automatically add citations in the answer.

### Requirements
Specifically, when a factual statement S in the answer uses information from context snippets l1, l2, ..., ln, please add citations by appending these snippet numbers to S in the format "{S}[{l1}][{l2}]...[{ln}]". For other sentences such as such as introductory sentences, summarization sentences, reasoning, and inference, you do not need to append anything to them to indicate they need no citations.
Citations should clearly show which facts are derived from the context snippets.
The section should contain tokens no more than 400.
The section should contain 5 citations at least.
Generate the section content without title.

### Input
[Topic]
The metallicity of PSR B1620 b and its uniqueness compared to other stars

[Section Title]
Comparison of PSR B1620 b’s Metallicity with Other Stars

[References Start]
[1]
PSR B1620 b’s metallicity is [Fe/H] = -2.09, which is 1% of the Sun’s metallicity. The measurement of stellar metallicity is crucial in understanding the relationship between the chemical composition of a star and its potential to host planetary systems. The low metallicity of PSR B1620 b places it among the most metal-poor stars known to host a planet.
[2]
A survey of 160 metal-poor main-sequence stars found no evidence for Jovian planets. This study is part of a larger effort to examine the correlation between a star’s metallicity and the types of planets it can host. The absence of giant planets in such stars adds to the rarity of PSR B1620 b’s exoplanet discovery.
[3]
Studies on the relationship between metallicity and planet formation suggest that lower metallicity may correlate with a lower probability of detecting giant planets. This is based on the hypothesis that higher metallicity facilitates the accumulation of material needed to form gas giants. As a result, metal-poor stars like PSR B1620 b are less likely to have Jupiter-like planets.
[4]
The detection of planets around such low-metallicity stars like PSR B1620 b challenges existing theories of planet formation, which typically emphasize the role of higher metallicity. According to the core-accretion model of planet formation, a certain threshold of metallicity is needed for the formation of large planetary cores that can then capture gas to form gas giants. PSR B1620 b’s planet contradicts this traditional view.
[5]
Research on exoplanet host stars has shown that metal-rich stars are significantly more likely to host gas giants, further emphasizing the unusual nature of PSR B1620 b’s exoplanet. Such findings have led scientists to reconsider the role of metallicity in planetary formation, suggesting that other factors might be at play, such as stellar age or specific evolutionary pathways.

### Output
PSR B1620 b has a metallicity of [Fe/H] = -2.09, which is about 1% of the Sun’s metallicity [1].Additionally, the survey of 160 metal-poor main-sequence stars found no evidence for Jovian planets, which highlights the uniqueness of PSR B1620 b’s metallicity [2].Studies on the relationship between metallicity and planet formation suggest that lower metallicity may correlate with a lower probability of detecting giant planets. This supports the observation that the majority of metal-poor stars do not host gas giants [3].The detection of planets around such low-metallicity stars like PSR B1620 b challenges existing theories of planet formation, which typically emphasize the role of higher metallicity in forming solid cores necessary for giant planet formation [4].Research on exoplanet host stars has shown that metal-rich stars are significantly more likely to host gas giants, further emphasizing the unusual nature of PSR B1620 b’s exoplanet. This has led to debates on whether metallicity alone can explain planet formation [5].

### Task instruction
Your task is to generate a section of an overall and structured literature review answer based on the provided reference and automatically add citations in the answer.

### Requirements
Specifically, when a factual statement S in the answer uses information from context snippets l1, l2, ..., ln, please add citations by appending these snippet numbers to S in the format "{S}[{l1}][{l2}]...[{ln}]". For other sentences such as such as introductory sentences, summarization sentences, reasoning, and inference, you do not need to append anything to them to indicate they need no citations.
Citations should clearly show which facts are derived from the context snippets.
The section should contain tokens no more than 400.
The section should contain 5 citations at least.
Generate the section content without title.

### Input
    """
    topic = dspy.InputField(prefix="[Topic]\n", format=str)
    section = dspy.InputField(prefix="[Section Title]\n", format=str)
    info = dspy.InputField(prefix="[References Start]\n",format=str)
    output = dspy.OutputField(prefix="\n[References End]\n\n### Output\n",format=str)